#
# Copyright (C) Since 2018 ~ 2025, Inc. All Rights Reserved
#
# @Time    : 2025/1/9-下午5:40
# @Author  : Timmy
# @Email   : tangping@cloudstream-tech.com
# @File    : ws.py
# @Software: lottery
# @Function:
#
import json
import os

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()
static_dir = os.path.join(os.path.dirname(__file__), "static")
app.mount("../static", StaticFiles(directory=static_dir), name="static")


class ConnectionManager:
    def __init__(self):
        self.frontend: list[WebSocket] = []
        self.terminal: list[WebSocket] = []
        self.active_connections: dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket, client_id: str) -> None:
        await websocket.accept()

        match client_id:
            case "frontend":
                self.frontend.append(websocket)
            case "terminal":
                self.terminal.append(websocket)
            case _:
                if client_id in self.active_connections:
                    msg = f"Client {client_id} already connected"
                    client = self.active_connections[client_id]
                    await client.send_text(msg)
                    await client.close()
                else:
                    self.active_connections[client_id] = websocket

    def disconnect(self, websocket: WebSocket, client_id: str) -> None:
        match client_id:
            case "frontend":
                self.frontend.remove(websocket)
            case "terminal":
                self.terminal.remove(websocket)
            case _:
                if client_id in self.active_connections:
                    del self.active_connections[client_id]

    async def send_personal_message(self, message: str, client_id: str) -> None:
        match client_id:
            case "frontend":
                for websocket in self.frontend:
                    await websocket.send_text(message)
            case "terminal":
                for websocket in self.terminal:
                    await websocket.send_text(message)
            case _:
                if client_id in self.active_connections:
                    await self.active_connections[client_id].send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections.values():
            await connection.send_text(message)
        for connection in self.frontend:
            await connection.send_text(message)

    async def send_to_frontend(self, message: str) -> None:
        await self.send_personal_message(message, client_id="frontend")

    async def send_to_terminal(self, message: str) -> None:
        await self.send_personal_message(message, client_id="terminal")


manager = ConnectionManager()


@app.get("/")
async def get():
    return RedirectResponse("/static/index.html")


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_text()
            match client_id:
                case "frontend":
                    await manager.send_to_terminal(data)
                case "terminal":
                    await manager.send_to_frontend(data)
                case _:
                    await manager.send_personal_message(f"You wrote: {data}", client_id)
                    await manager.broadcast(f"Client #{client_id} says: {data}")

    except WebSocketDisconnect:
        manager.disconnect(websocket, client_id)
        # await manager.broadcast(f"Client #{client_id} left the chat")


async def websocket_server():
    config = uvicorn.Config(app, host='0.0.0.0', port=8000, log_level="info", workers=1)
    server = uvicorn.Server(config)
    await server.serve()


def websocket_server_sync():
    import asyncio
    try:
        asyncio.run(websocket_server())
    except KeyboardInterrupt:
        print("Keyboard interrupt received, shutting down")


if __name__ == '__main__':
    websocket_server_sync()
